from input import Input

from plugin_system import Plugin

import random

import numpy as np

class RandomPlanetsGenerator(Input):
    def __init__(self, xyz_filename=None, params_filename=None, num_planets=None, seed=None):
        self.xyz_flname = xyz_filename
        self.xyz_fl = None
        self.params_flname = params_filename
        self.params_fl = None
        self.num_planets = num_planets
        if seed is not -1:
            random.seed(seed)
        self.state = random.getstate()
        
    def initialize(self):
        self.xyz_fl = open(self.xyz_flname)
        self.params_fl = open(self.params_flname)

    def finalize(self):
        self.xyz_fl.close()
        self.params_fl.close()

    def read(self, system):
        masses = []
        mass_g = {}
        for ln in self.params_fl:
            cols = ln.split()
            name = cols[0]
            mass_g[name] = np.float64(cols[1])
        G = mass_g["Sun"]
        for name, mass in mass_g.iteritems():
            masses.append(mass / G)
        max_mass = max(masses)
        min_mass = min(masses)
        
        all_positions = []
        all_velocities = []
        ln = self.xyz_fl.readline()
        while ln != "":
            num_objects, step = map(int, ln.split())
            for j in xrange(num_objects):
                cols = self.xyz_fl.readline().split()
                positions = np.array(map(float, cols[1:4]), np.float64)
                velocities = np.array(map(float, cols[4:7]), np.float64)
                all_positions.append(positions)
                all_velocities.append(velocities)
            ln = self.xyz_fl.readline()
        
        x_pos, y_pos, z_pos = zip(*all_positions)
        x_vel, y_vel, z_vel = zip(*all_velocities)
        
        x_pos_range = min(x_pos), max(x_pos)
        y_pos_range = min(y_pos), max(y_pos)
        z_pos_range = min(z_pos), max(z_pos)

        x_vel_range = min(x_vel), max(x_vel)
        y_vel_range = min(y_vel), max(y_vel)
        z_vel_range = min(z_vel), max(z_vel)

        random.setstate(self.state)
        for i in xrange(self.num_planets):
            name = "random_%s" % (i + 1)
            mass = random.uniform(min_mass, max_mass)
            position = np.array([random.uniform(*x_pos_range), random.uniform(*y_pos_range), 
                        random.uniform(*z_pos_range)], np.float64)
            velocity = np.array([random.uniform(*x_vel_range), random.uniform(*y_vel_range), 
                        random.uniform(*z_vel_range)], np.float64)
            system.planet_names.append(name)
            system.masses[name] = mass
            system.positions[name] = position
            system.velocities[name] = velocity


class RandomPlanetsGeneratorPluginAdaptor(Plugin, RandomPlanetsGenerator):
    plugin_name = "RandomPlanetsGenerator"
    parameter_definitions = {"xyz_filename" : (str, True, "Positions and velocities file"),
                             "params_filename" : (str, True, "Parameters file"),
                             "num_planets" : (int, True, "Number of random planets to generate"),
                             "seed" : (int, False, "Seed for random number generator.", -1)}

    def __init__(self, *args, **user_params):
        params = self.configure_parameters(**user_params)
        RandomPlanetsGenerator.__init__(self, xyz_filename=params["xyz_filename"], 
                                        params_filename=params["params_filename"],
                                        num_planets=params["num_planets"],
                                        seed=params["seed"])
